# encoding: utf-8
"""
调用即梦API生成视频
"""
from pydantic import BaseModel
import logging
import time
import requests
import uuid
import os
import datetime
import random
from . import utils
from .file_service import upload_file
from typing import Dict

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

API_TOKEN = "7433e85ddc99ed88aa8656aacf6630dc"  # 你登录即梦获得的session_id

# 在读取 video_output_path 后添加目录检查和创建逻辑
video_output_path = '/videos'


class VideoRequest(BaseModel):
    prompt: str
    aspect_ratio: str = "16:9"
    duration_ms: int = 5000
    fps: int = 24


# 添加新的辅助函数
def generate_timestamp_filename_for_video(extension='mp4'):
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    random_number = random.randint(1000, 9999)
    filename = f"video_{timestamp}_{random_number}.{extension}"
    return filename


def download_video(url, output_path):
    response = requests.get(url, stream=True)
    response.raise_for_status()

    filename = generate_timestamp_filename_for_video()
    file_path = os.path.join(output_path, filename)

    with open(file_path, 'wb') as file:
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                file.write(chunk)

    return filename, file_path


def generate_cookie(token: str) -> str:
    """生成Cookie
    """
    return f"sessionid={token}; sessionid_ss={token}; sid_tt={token}; uid_tt={token}; uid_tt_ss={token}"


async def generate_video(request: VideoRequest) -> Dict:
    try:
        logger.info(f"generate_video API 调用开始，提示词: {request.prompt}")
        start_time = time.time()

        token = API_TOKEN
        device_time = utils.get_timestamp()
        uri = "/mweb/v1/generate_video?aid=513695"
        VERSION_CODE = "5.8.0"
        PLATFORM_CODE = "7"
        sign = utils.md5(f"9e2c|{uri[-7:]}|{PLATFORM_CODE}|{VERSION_CODE}|{device_time}||11ac")

        # 初始化视频生成API相关配置
        video_api_headers = {
            'accept': 'application/json, text/plain, */*',
            'accept-language': 'zh-CN,zh;q=0.9',
            'app-sdk-version': '48.0.0',
            'appid': '513695',
            'appvr': '5.8.0',
            'content-type': 'application/json',
            'cookie': generate_cookie(token),
            'device-time': str(int(time.time())),
            'lan': 'zh-Hans',
            'loc': 'cn',
            'origin': 'https://jimeng.jianying.com',
            'pf': '7',
            'priority': 'u=1, i',
            'referer': 'https://jimeng.jianying.com/ai-tool/video/generate',
            'sec-ch-ua': '"Google Chrome";v="129", "Not=A?Brand";v="8", "Chromium";v="129"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Windows"',
            'sec-fetch-dest': 'empty',
            'sec-fetch-mode': 'cors',
            'sec-fetch-site': 'same-origin',
            'sign': sign,
            'sign-ver': '1',
            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36'
        }

        video_api_base = "https://jimeng.jianying.com/mweb/v1"

        # 生成唯一的submit_id
        submit_id = str(uuid.uuid4())

        # 准备请求数据
        generate_video_payload = {
            "submit_id": submit_id,
            "task_extra": "{\"promptSource\":\"custom\",\"originSubmitId\":\"0340110f-5a94-42a9-b737-f4518f90361f\",\"isDefaultSeed\":1,\"originTemplateId\":\"\",\"imageNameMapping\":{},\"isUseAiGenPrompt\":false,\"batchNumber\":1}",
            "http_common_info": {"aid": 513695},
            "input": {
                "video_aspect_ratio": request.aspect_ratio,
                "seed": 2934141961,
                "video_gen_inputs": [
                    {
                        "prompt": request.prompt,
                        "fps": request.fps,
                        "duration_ms": request.duration_ms,
                        "video_mode": 2,
                        "template_id": ""
                    }
                ],
                "priority": 0,
                "model_req_key": "dreamina_ic_generate_video_model_vgfm_lite"
            },
            "mode": "workbench",
            "history_option": {},
            "commerce_info": {
                "resource_id": "generate_video",
                "resource_id_type": "str",
                "resource_sub_type": "aigc",
                "benefit_type": "basic_video_operation_vgfm_lite"
            },
            "client_trace_data": {}
        }

        # 发送生成视频请求
        generate_video_url = f"{video_api_base}/generate_video?aid=513695"
        logger.info(f"发送视频生成请求...")

        response = requests.post(generate_video_url,
                                 headers=video_api_headers,
                                 json=generate_video_payload
                                 )

        if response.status_code != 200:
            return {
                "success": False,
                "error": f"视频生成请求失败，状态码：{response.status_code}",
                "video": None
            }

        response_data = response.json()
        if not response_data or "data" not in response_data or "aigc_data" not in response_data["data"]:
            return {
                "success": False,
                "error": '视频生成接口返回格式错误',
                "video": None
            }

        task_id = response_data["data"]["aigc_data"]["task"]["task_id"]
        logger.info(f"视频生成任务已创建，任务ID: {task_id}")

        # 轮询检查视频生成状态
        mget_generate_task_url = f"{video_api_base}/mget_generate_task?aid=513695"
        mget_generate_task_payload = {"task_id_list": [task_id]}

        # 最多尝试30次，每次间隔2秒
        for attempt in range(50):
            time.sleep(2)
            logger.info(f"检查视频状态，第 {attempt + 1} 次尝试...")

            response = requests.post(mget_generate_task_url,
                                     headers=video_api_headers,
                                     json=mget_generate_task_payload)
            if response.status_code != 200:
                logger.warning(f"状态检查失败，状态码：{response.status_code}")
                continue

            response_data = response.json()
            if not response_data or "data" not in response_data or "task_map" not in response_data["data"]:
                logger.warning("状态检查返回格式错误")
                continue

            task_data = response_data["data"]["task_map"].get(task_id)
            if not task_data:
                logger.warning(f"未找到任务 {task_id} 的状态信息")
                continue

            task_status = task_data.get("status")
            logger.info(f"任务状态: {task_status}")

            if task_status == 50:  # 视频生成完成
                if "item_list" in task_data and task_data["item_list"] and "video" in task_data["item_list"][0]:
                    video_data = task_data["item_list"][0]["video"]
                    if "transcoded_video" in video_data and "origin" in video_data["transcoded_video"]:
                        video_url = video_data["transcoded_video"]["origin"]["video_url"]
                        elapsed_time = time.time() - start_time
                        logger.info(f"视频生成成功，耗时 {elapsed_time:.2f} 秒，URL: {video_url}")

                        # 下载视频到本地
                        try:
                            filename, file_path = download_video(video_url, video_output_path)
                            logger.info(f"视频已下载到本地: {file_path}")
                            file_url = upload_file(file_path)
                            return {
                                "success": True,
                                "error": None,
                                "video": file_url
                            }
                        except Exception as e:
                            logger.error(f"处理视频文件失败: {str(e)}")
                            return {
                                "success": False,
                                "error": str(e),
                                "video": None
                            }

        return {
            "success": False,
            "error": '视频生成超时',
            "video": None
        }

    except Exception as e:
        logger.error(f"视频生成失败: {str(e)}")
        return {
            "success": False,
            "error": str(e),
            "video": None
        }
